#!/usr/bin/python
#--------------------------------------------#
# pyclaes.py
#--------------------------------------------#

import argparse
import numpy as np

from pyclaes_ecb import claes_ecb


# TODO: intercept exceptions here
def main():
    parser = argparse.ArgumentParser(description="Encrypt files using AES algorithm")
    parser.add_argument('--aes-ecb', action='store_true')
    parser.add_argument('--profiling', action='store_true')
    parser.add_argument('--overlap', action='store_true')
    parser.add_argument('--use-cpu', action='store_true')
    parser.add_argument('key')
    parser.add_argument('inFile')
    parser.add_argument('outFile')
    args = parser.parse_args()
    
    # Init key and data arrays
    #key = np.fromstring(args.key, dtype=np.uint8)
    #data = np.fromstring(open(args.inFile, 'rb').read(), dtype=np.uint8)
    data = open(args.inFile, 'rb').read()

    # Pad data with zeros if necessary (size must be a multiple of 16)
    if len(data)&0x0f:
        data += '\0'*(16-(len(data)&0x0f))

    # TODO: has new modes of operation get added, handle them
    enc = claes_ecb(args.profiling, args.use_cpu)
    cipherText = enc.encrypt(args.key, data, overlap=args.overlap)

    # Write back to disk cipher text
    open(args.outFile, 'wb').write(cipherText)

    return 0


if __name__=='__main__':
    exit(main())
